{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.transforms import RandAugment\n",
    "from torch.optim import lr_scheduler\n",
    "import numpy as np\n",
    "import time\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "from Resnet18 import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_data_loader():    \n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        RandAugment(),  # 使用 RandAugment\n",
    "        transforms.ToTensor(),\n",
    "        #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "        transforms.Normalize(\n",
    "            (0.5070751592371323, 0.48654887331495095, 0.4409178433670343),\n",
    "            (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)\n",
    "        )\n",
    "    ])\n",
    "    \n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "    ])\n",
    "    \n",
    "\n",
    "    trainset = torchvision.datasets.CIFAR100(root='./data', train=True,\n",
    "                                             download=False, transform=transform_train)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,\n",
    "                                              shuffle=True, num_workers=2)\n",
    "\n",
    "    testset = torchvision.datasets.CIFAR100(root='./data', train=False,\n",
    "                                            download=False, transform=transform_test)\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=128,\n",
    "                                             shuffle=False, num_workers=2)\n",
    "        # 实例化模型\n",
    "    return trainloader,testloader\n",
    "\n",
    "\n",
    "def plot_training_results(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs):\n",
    "    # Convert PyTorch tensors to NumPy arrays\n",
    "    train_losses = np.array(train_losses)\n",
    "    train_top1_accs = np.array(train_top1_accs)\n",
    "    train_top5_accs = np.array(train_top5_accs)\n",
    "    val_losses = np.array(val_losses)\n",
    "    val_top1_accs = np.array(val_top1_accs)\n",
    "    val_top5_accs = np.array(val_top5_accs)\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    \n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(train_losses, label='训练损失')\n",
    "    plt.plot(val_losses, label='验证损失')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('训练/验证损失曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(train_top1_accs, label='训练 Top-1 准确率')\n",
    "    plt.plot(val_top1_accs, label='验证 Top-1 准确率')\n",
    "    plt.plot(train_top5_accs, label='训练 Top-5 准确率')\n",
    "    plt.plot(val_top5_accs, label='验证 Top-5 准确率')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('训练/验证准确率曲线')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# 训练模型\n",
    "def train_model(model,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "       \n",
    "    train_losses = []\n",
    "    train_top1_accs = []\n",
    "    train_top5_accs = []\n",
    "    val_losses = []\n",
    "    val_top1_accs = []\n",
    "    val_top5_accs = []\n",
    "\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    " \n",
    "    for epoch in range(num_epochs):\n",
    "        print('第 {}/{} 轮训练'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # 每个epoch都有训练和验证阶段\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # 设置模型为训练模式\n",
    "            else:\n",
    "                model.eval()   # 设置模型为评估模式\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "            running_top1_corrects = 0\n",
    "            running_top5_corrects = 0\n",
    "\n",
    "            # 遍历数据\n",
    "            for inputs, labels in (trainloader if phase == 'train' else testloader):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # 零化参数梯度\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # 前向传播\n",
    "                # 只在训练阶段追踪历史\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    loss = criterion(outputs, labels)\n",
    "\n",
    "                    # 计算 top-1 和 top-5 准确率\n",
    "                    _, top5_preds = torch.topk(outputs, 5, dim=1)\n",
    "                    top1_correct = torch.sum(preds == labels.data)\n",
    "                    top5_correct = torch.sum(top5_preds == labels.view(-1, 1))\n",
    "\n",
    "                    # 只有在训练阶段反向传播+优化\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # 统计\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += top1_correct\n",
    "                running_top1_corrects += top1_correct\n",
    "                running_top5_corrects += top5_correct\n",
    "\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top1_acc = running_top1_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top5_acc = running_top5_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "\n",
    "            print('{} 损失: {:.4f} Top-1 准确率: {:.4f} Top-5 准确率: {:.4f}'.format(\n",
    "                phase, epoch_loss, epoch_top1_acc, epoch_top5_acc))\n",
    "            \n",
    "            # 保存损失和准确率\n",
    "            if phase == 'train':\n",
    "                train_losses.append(epoch_loss)\n",
    "                train_top1_accs.append(epoch_top1_acc)\n",
    "                train_top5_accs.append(epoch_top5_acc)\n",
    "            else:\n",
    "                val_losses.append(epoch_loss)\n",
    "                val_top1_accs.append(epoch_top1_acc)\n",
    "                val_top5_accs.append(epoch_top5_acc)\n",
    "\n",
    "            # 深度复制模型\n",
    "            if phase == 'val' and epoch_top1_acc > best_acc:\n",
    "                best_acc = epoch_top1_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('训练完成，耗时 {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('最佳验证集准确率: {:4f}'.format(best_acc))\n",
    "\n",
    "    # 加载最佳模型权重\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    \n",
    "    # 绘制训练结果曲线\n",
    "\n",
    "    return model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n",
      "第 0/49 轮训练\n",
      "----------\n",
      "train 损失: 3.9639 Top-1 准确率: 0.0926 Top-5 准确率: 0.2857\n",
      "val 损失: 4.0318 Top-1 准确率: 0.0958 Top-5 准确率: 0.3030\n",
      "\n",
      "第 1/49 轮训练\n",
      "----------\n",
      "train 损失: 3.2084 Top-1 准确率: 0.2125 Top-5 准确率: 0.4998\n",
      "val 损失: 3.3865 Top-1 准确率: 0.2156 Top-5 准确率: 0.5171\n",
      "\n",
      "第 2/49 轮训练\n",
      "----------\n",
      "train 损失: 2.7141 Top-1 准确率: 0.3026 Top-5 准确率: 0.6241\n",
      "val 损失: 3.4521 Top-1 准确率: 0.2152 Top-5 准确率: 0.5262\n",
      "\n",
      "第 3/49 轮训练\n",
      "----------\n",
      "train 损失: 2.3983 Top-1 准确率: 0.3719 Top-5 准确率: 0.6946\n",
      "val 损失: 3.0113 Top-1 准确率: 0.3070 Top-5 准确率: 0.6051\n",
      "\n",
      "第 4/49 轮训练\n",
      "----------\n",
      "train 损失: 2.1671 Top-1 准确率: 0.4197 Top-5 准确率: 0.7444\n",
      "val 损失: 2.3000 Top-1 准确率: 0.3992 Top-5 准确率: 0.7379\n",
      "\n",
      "第 5/49 轮训练\n",
      "----------\n",
      "train 损失: 1.9814 Top-1 准确率: 0.4622 Top-5 准确率: 0.7791\n",
      "val 损失: 2.2223 Top-1 准确率: 0.4225 Top-5 准确率: 0.7418\n",
      "\n",
      "第 6/49 轮训练\n",
      "----------\n",
      "train 损失: 1.8327 Top-1 准确率: 0.4976 Top-5 准确率: 0.8060\n",
      "val 损失: 2.1312 Top-1 准确率: 0.4346 Top-5 准确率: 0.7472\n",
      "\n",
      "第 7/49 轮训练\n",
      "----------\n",
      "train 损失: 1.7242 Top-1 准确率: 0.5232 Top-5 准确率: 0.8266\n",
      "val 损失: 2.1326 Top-1 准确率: 0.4613 Top-5 准确率: 0.7592\n",
      "\n",
      "第 8/49 轮训练\n",
      "----------\n",
      "train 损失: 1.6203 Top-1 准确率: 0.5487 Top-5 准确率: 0.8431\n",
      "val 损失: 2.1391 Top-1 准确率: 0.4563 Top-5 准确率: 0.7740\n",
      "\n",
      "第 9/49 轮训练\n",
      "----------\n",
      "train 损失: 1.5462 Top-1 准确率: 0.5676 Top-5 准确率: 0.8552\n",
      "val 损失: 1.9393 Top-1 准确率: 0.4840 Top-5 准确率: 0.7992\n",
      "\n",
      "第 10/49 轮训练\n",
      "----------\n",
      "train 损失: 1.4618 Top-1 准确率: 0.5914 Top-5 准确率: 0.8678\n",
      "val 损失: 1.8283 Top-1 准确率: 0.5198 Top-5 准确率: 0.8108\n",
      "\n",
      "第 11/49 轮训练\n",
      "----------\n",
      "train 损失: 1.3930 Top-1 准确率: 0.6078 Top-5 准确率: 0.8772\n",
      "val 损失: 1.9885 Top-1 准确率: 0.4857 Top-5 准确率: 0.7798\n",
      "\n",
      "第 12/49 轮训练\n",
      "----------\n",
      "train 损失: 1.3374 Top-1 准确率: 0.6214 Top-5 准确率: 0.8860\n",
      "val 损失: 1.7277 Top-1 准确率: 0.5466 Top-5 准确率: 0.8276\n",
      "\n",
      "第 13/49 轮训练\n",
      "----------\n",
      "train 损失: 1.2807 Top-1 准确率: 0.6358 Top-5 准确率: 0.8955\n",
      "val 损失: 1.6229 Top-1 准确率: 0.5680 Top-5 准确率: 0.8442\n",
      "\n",
      "第 14/49 轮训练\n",
      "----------\n",
      "train 损失: 1.2385 Top-1 准确率: 0.6473 Top-5 准确率: 0.8994\n",
      "val 损失: 1.7065 Top-1 准确率: 0.5476 Top-5 准确率: 0.8319\n",
      "\n",
      "第 15/49 轮训练\n",
      "----------\n",
      "train 损失: 1.1914 Top-1 准确率: 0.6590 Top-5 准确率: 0.9067\n",
      "val 损失: 1.7647 Top-1 准确率: 0.5463 Top-5 准确率: 0.8231\n",
      "\n",
      "第 16/49 轮训练\n",
      "----------\n",
      "train 损失: 1.1365 Top-1 准确率: 0.6723 Top-5 准确率: 0.9136\n",
      "val 损失: 1.7623 Top-1 准确率: 0.5471 Top-5 准确率: 0.8194\n",
      "\n",
      "第 17/49 轮训练\n",
      "----------\n",
      "train 损失: 1.1048 Top-1 准确率: 0.6808 Top-5 准确率: 0.9174\n",
      "val 损失: 1.5234 Top-1 准确率: 0.5893 Top-5 准确率: 0.8642\n",
      "\n",
      "第 18/49 轮训练\n",
      "----------\n",
      "train 损失: 1.0609 Top-1 准确率: 0.6933 Top-5 准确率: 0.9226\n",
      "val 损失: 1.4350 Top-1 准确率: 0.6084 Top-5 准确率: 0.8712\n",
      "\n",
      "第 19/49 轮训练\n",
      "----------\n",
      "train 损失: 1.0256 Top-1 准确率: 0.7001 Top-5 准确率: 0.9283\n",
      "val 损失: 1.7365 Top-1 准确率: 0.5600 Top-5 准确率: 0.8465\n",
      "\n",
      "第 20/49 轮训练\n",
      "----------\n",
      "train 损失: 1.0006 Top-1 准确率: 0.7086 Top-5 准确率: 0.9317\n",
      "val 损失: 1.4748 Top-1 准确率: 0.5988 Top-5 准确率: 0.8745\n",
      "\n",
      "第 21/49 轮训练\n",
      "----------\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[10], line 11\u001b[0m\n\u001b[0;32m      9\u001b[0m optimizer \u001b[38;5;241m=\u001b[39m optim\u001b[38;5;241m.\u001b[39mSGD(net\u001b[38;5;241m.\u001b[39mparameters(), lr\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.01\u001b[39m, momentum\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.9\u001b[39m, weight_decay\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m5e-4\u001b[39m)\n\u001b[0;32m     10\u001b[0m scheduler \u001b[38;5;241m=\u001b[39m lr_scheduler\u001b[38;5;241m.\u001b[39mStepLR(optimizer, step_size\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m50\u001b[39m, gamma\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m0.1\u001b[39m)\n\u001b[1;32m---> 11\u001b[0m model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs) \u001b[38;5;241m=\u001b[39m \u001b[43mtrain_model\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnet\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtrainloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43mtestloader\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcriterion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptimizer\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mscheduler\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnum_epochs\u001b[49m\u001b[43m)\u001b[49m  \n",
      "Cell \u001b[1;32mIn[9], line 125\u001b[0m, in \u001b[0;36mtrain_model\u001b[1;34m(model, trainloader, device, testloader, criterion, optimizer, scheduler, num_epochs)\u001b[0m\n\u001b[0;32m    122\u001b[0m         optimizer\u001b[38;5;241m.\u001b[39mstep()\n\u001b[0;32m    124\u001b[0m \u001b[38;5;66;03m# 统计\u001b[39;00m\n\u001b[1;32m--> 125\u001b[0m running_loss \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[43mloss\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mitem\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m*\u001b[39m inputs\u001b[38;5;241m.\u001b[39msize(\u001b[38;5;241m0\u001b[39m)\n\u001b[0;32m    126\u001b[0m running_corrects \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m top1_correct\n\u001b[0;32m    127\u001b[0m running_top1_corrects \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m top1_correct\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    " ## 修改模型\n",
    "net = ResNet18().to(device)\n",
    "num_epochs=50\n",
    " # 定义损失函数和优化器\n",
    "trainloader,testloader = get_data_loader()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)\n",
    "model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs) = train_model(net,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs)  \n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_cpu_numpy(tensor):\n",
    "    return tensor.cpu().numpy()\n",
    "def to_cpu(list_):\n",
    "    if isinstance(list_, list):\n",
    "        return list(map(to_cpu_numpy,list_))\n",
    "    else:\n",
    "        return to_cpu_numpy(list_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'train_losses' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[12], line 1\u001b[0m\n\u001b[1;32m----> 1\u001b[0m plot_training_results(\u001b[43mtrain_losses\u001b[49m, to_cpu(train_top1_accs), to_cpu(train_top5_accs), val_losses, to_cpu(val_top1_accs), to_cpu(val_top5_accs))\n",
      "\u001b[1;31mNameError\u001b[0m: name 'train_losses' is not defined"
     ]
    }
   ],
   "source": [
    "plot_training_results(train_losses, to_cpu(train_top1_accs), to_cpu(train_top5_accs), val_losses, to_cpu(val_top1_accs), to_cpu(val_top5_accs))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
